package com.cy.duosearch.langchain.agent;

import com.cy.duosearch.langchain.tool.*;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiChatRequestParameters;
import io.micrometer.common.util.StringUtils;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;

@AllArgsConstructor
@Data
@Builder
@Slf4j
public class ToolCallAgent extends ReactAgent {


    private ChatMemory memory;
    private ChatLanguageModel languageModel;
    private String prompt;

    List<ToolExecutionRequest> toolExecutionRequests = new ArrayList<>();

    private static final String SYSTEM_PROMPT = """
        你是一个全能的 AI 助手，可以解决用户提出的任何任务。你可以调用各种工具来完成各种复杂的请求。
        你可以进行编程，浏览网页，进行网页信息检索，进行处理等等。
        """;

    private static final String STEP_PROMPT = """
        你可以使用以下工具与计算机交互:
        - SearchTool：执行网页信息检索
        - FileSaverTool：本地保存文件，例如 txt、py、html 等
        - TerminateTool：使用此工具可以表明已完成请求
        根据用户需求，主动选择最合适的工具或工具组合。对于复杂任务，可以将问题拆解，并逐步使用不同工具来解决。在使用每个工具后，清楚地解释执行结果，并建议下一步行动
        """;



//    """
//        你可以使用以下工具与计算机交互：
//        - PythonExecute：执行 Python 代码，与计算机系统交互，进行数据处理
//        - FileSaverTool：本地保存文件，例如 txt、py、html 等
//        - BowerScreenshotTool：使用浏览器进行截屏和截图
//        - WebCrawlingTool：对网页的内容进行解析和提取文本
//        - SearchTool：执行网页信息检索
//        - TerminateTool：使用此工具可以表明已完成请求
//        根据用户需求，主动选择最合适的工具或工具组合。对于复杂任务，可以将问题拆解，并逐步使用不同工具来解决。在使用每个工具后，清楚地解释执行结果，并建议下一步行动
//        """;

    @Override
    public boolean think() {

        // 首次需要把系统prompt添加到chat中
        if (1 == getCurrentStep()) {
            getMemory().add(SystemMessage.systemMessage(SYSTEM_PROMPT));
            getMemory().add(UserMessage.userMessage(getPrompt() + STEP_PROMPT));
        } else {
            getMemory().add(UserMessage.userMessage(STEP_PROMPT));
        }


        //将自定义工具转化为要求的工具类
        List<ToolSpecification> toolSpecificationsList = getTools().stream()
                .map(
                        v -> v.getToolSpecification()
                ).toList();

        ChatResponse chatResponse;

        if (getLanguageModel() instanceof OpenAiChatModel) {
            OpenAiChatModel openAiChatModel = (OpenAiChatModel) getLanguageModel();
            OpenAiChatRequestParameters parameters = OpenAiChatRequestParameters.builder()
                    .toolSpecifications(toolSpecificationsList)
                    .modelName("QwQ-32B")
                    .build();
            chatResponse = openAiChatModel.doChat(ChatRequest.builder()
                    .messages(getMemory().messages())
                    .parameters(parameters)
                    .build());
        } else {
            // 使用通用的 doChat 方法
            chatResponse = getLanguageModel().doChat(ChatRequest.builder()
                    .messages(getMemory().messages())
                    .toolSpecifications(toolSpecificationsList)
                    .build());
        }



        if (StringUtils.isNotEmpty(chatResponse.aiMessage().text())) {
            log.info("ai thinking {}", chatResponse.aiMessage().text());

        }

        if (chatResponse.aiMessage().hasToolExecutionRequests()) {
            toolExecutionRequests = chatResponse.aiMessage().toolExecutionRequests();
//            log.info("choose {} tool execution requests", chatResponse.aiMessage().toolExecutionRequests().size());
        }
        return chatResponse.aiMessage().hasToolExecutionRequests();
    }

    @Override
    public ToolResult action() {

        ToolResult toolResult = null;
        // ai 会把 函数名和函数相关参数给我们，我们 再 通过反射去调用
        // do actual job

        for (ToolExecutionRequest toolExecutionRequest : toolExecutionRequests) {
            log.info("当前执行工具的名称：{}", toolExecutionRequest.name());

            toolResult = ToolUtil.runTool(toolExecutionRequest);
            String result = "Observed output of cmd '%s' , executed : %s ".formatted(toolExecutionRequest.name(),
                    toolResult.getResult());

            getMemory().add(AiMessage.aiMessage(result));  //react模式 把回答继续添加到对话，也会影响后续的思考
        }

        return toolResult;
    }


}
